import os
import io
import PyPDF2
from docx import Document
from typing import Dict, Any, Optional

class DocumentParser:
    def __init__(self):
        """
        初始化文档解析器，支持PDF和Word文档解析
        """
        pass

    def parse_document(self, file_path: str, file_content: bytes = None) -> Dict[str, Any]:
        """
        解析文档内容
        :param file_path: 文件路径
        :param file_content: 文件内容（字节流），如果提供则优先使用
        :return: 解析结果字典，包含文本内容和元数据
        """
        # 获取文件扩展名
        _, ext = os.path.splitext(file_path)
        ext = ext.lower()

        # 根据文件类型选择不同的解析方法
        if ext == '.pdf':
            return self._parse_pdf(file_path, file_content)
        elif ext == '.docx':
            return self._parse_docx(file_path, file_content)
        elif ext == '.txt':
            return self._parse_txt(file_path, file_content)
        else:
            raise ValueError(f"不支持的文件类型: {ext}")

    def _parse_pdf(self, file_path: str, file_content: Optional[bytes] = None) -> Dict[str, Any]:
        """
        解析PDF文件
        :param file_path: PDF文件路径
        :param file_content: PDF文件内容（字节流）
        :return: 解析结果字典
        """
        text = ""
        metadata = {}

        try:
            if file_content:
                # 使用提供的文件内容，包装在BytesIO中
                pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_content))
            else:
                # 从文件路径读取
                with open(file_path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)

            # 提取文本
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text += page.extract_text() or ""
                text += "\n"

            # 提取元数据
            if pdf_reader.metadata:
                metadata['title'] = pdf_reader.metadata.get('/Title', '').strip()
                metadata['author'] = pdf_reader.metadata.get('/Author', '').strip()
                metadata['creator'] = pdf_reader.metadata.get('/Creator', '').strip()
                metadata['producer'] = pdf_reader.metadata.get('/Producer', '').strip()
                metadata['creation_date'] = pdf_reader.metadata.get('/CreationDate', '').strip()

            return {
                'text': text,
                'metadata': metadata,
                'page_count': len(pdf_reader.pages)
            }
        except Exception as e:
            raise Exception(f"解析PDF文件时出错: {str(e)}")

    def _parse_docx(self, file_path: str, file_content: Optional[bytes] = None) -> Dict[str, Any]:
        """
        解析Word文档
        :param file_path: Word文档路径
        :param file_content: Word文档内容（字节流）
        :return: 解析结果字典
        """
        text = ""
        metadata = {}

        try:
            if file_content:
                # 使用提供的文件内容
                doc = Document(file_content)
            else:
                # 从文件路径读取
                doc = Document(file_path)

            # 提取文本
            for para in doc.paragraphs:
                text += para.text
                text += "\n"

            # 提取元数据
            core_properties = doc.core_properties
            metadata['title'] = core_properties.title or ''
            metadata['author'] = core_properties.author or ''
            metadata['created'] = str(core_properties.created) if core_properties.created else ''
            metadata['modified'] = str(core_properties.modified) if core_properties.modified else ''
            metadata['subject'] = core_properties.subject or ''
            metadata['keywords'] = core_properties.keywords or ''

            return {
                'text': text,
                'metadata': metadata
            }
        except Exception as e:
            raise Exception(f"解析Word文档时出错: {str(e)}")

    def _parse_txt(self, file_path: str, file_content: Optional[bytes] = None) -> Dict[str, Any]:
        """
        解析TXT文件，支持多种编码
        :param file_path: TXT文件路径
        :param file_content: TXT文件内容（字节流）
        :return: 解析结果字典
        """
        try:
            if file_content:
                # 尝试使用UTF-8解码
                try:
                    text = file_content.decode('utf-8')
                except UnicodeDecodeError:
                    # 尝试使用GBK解码
                    try:
                        text = file_content.decode('gbk')
                    except UnicodeDecodeError:
                        # 尝试使用GB2312解码
                        text = file_content.decode('gb2312')
            else:
                # 从文件路径读取，尝试多种编码
                encodings = ['utf-8', 'gbk', 'gb2312']
                for encoding in encodings:
                    try:
                        with open(file_path, 'r', encoding=encoding) as f:
                            text = f.read()
                        break
                    except UnicodeDecodeError:
                        continue
                else:
                    raise Exception(f"无法使用以下编码解析文件: {', '.join(encodings)}")

            return {
                'text': text,
                'metadata': {}
            }
        except Exception as e:
            raise Exception(f"解析TXT文件时出错: {str(e)}")